import os
from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_models import ChatZhipuAI 
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationChain

_ = load_dotenv(find_dotenv())

api_key = os.environ.get('ZHIPU_API_KEY')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")

model = os.environ.get('model')
if model is None:
    raise ValueError("model is not set in the .env file")

llm = ChatZhipuAI(api_key=api_key,
                  model=model,
                  temperature=0.9,              
                )

# 设定使用最后k个交互
memory=ConversationBufferMemory(human_prefix="User", ai_prefix="Assistant",)
conversation = ConversationChain(
    llm=llm, 
    verbose=True, 
    memory=memory
)
conversation.predict(input="我昨天给张三买了五本书，分别涉及编程、哲学、美食等主题。")
conversation.predict(input="我昨天给谁买书了？")
conversation.predict(input="我昨天买了几本书？")
conversation.predict(input="我昨天买的几本书涉及哪些主题？")

#print(f"\nConversationBufferMemory:\n{memory}")